/*
 * Copyright (c) <2025> <shihaorui> and <2501060546@qq.com>. All rights reserved.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Contributor: <shihaorui> (<2501060546@qq.com>)
 */
 #include "kernel_operator.h"
 using namespace AscendC;
 constexpr int32_t BUFFER_NUM = 2;
 
 template <typename T>
 struct SegmentMinTraits;
 
 template <>
 struct SegmentMinTraits<half> {
     static __aicore__ inline half InitValue() { return (half)65504.0f; }
     static __aicore__ inline half ZeroValue() { return (half)0.0f; }
     static __aicore__ inline bool IsClose(half a, half b) {
         float fa = static_cast<float>(a);
         float fb = static_cast<float>(b);
         float diff = fa > fb ? fa - fb : fb - fa;
         return (fa == fb) || (diff < 1e-4f);
     }
 };
 
 template <>
 struct SegmentMinTraits<float> {
     static __aicore__ inline float InitValue() { return 3.402823466e+38F; }
     static __aicore__ inline float ZeroValue() { return 0.0f; }
     static __aicore__ inline bool IsClose(float a, float b) {
         float diff = a > b ? a - b : b - a;
         return (a == b) || (diff < 1e-6f);
     }
 };
 
 template <>
 struct SegmentMinTraits<int32_t> {
     static __aicore__ inline int32_t InitValue() { return 0x7fffffff; }
     static __aicore__ inline int32_t ZeroValue() { return 0; }
     static __aicore__ inline bool IsClose(int32_t a, int32_t b) {
         return a == b;
     }
 };
 
 using SrcDataType = DTYPE_SRC;
 
 template <typename T>
 class SegmentMinCooKernel {
 public:
     __aicore__ inline SegmentMinCooKernel(){}
     __aicore__ inline void Init(GM_ADDR src, GM_ADDR index, GM_ADDR out, GM_ADDR workspace,
                                 uint32_t srcLength, uint32_t ubPartDataNum, 
                                 uint32_t E_1, uint32_t K, uint32_t coreDataNum, uint32_t coreTailDataNum,
                                 uint32_t nSegments, uint32_t ALIGN_NUM, uint32_t indexPhysicalNum, uint32_t aivNum,
                                 uint32_t logicalM, uint32_t logicalIndexNum, uint32_t strideE1, uint32_t strideM,
                                 uint32_t idxDimNum, const uint32_t* logicalShape, const uint32_t* idxPhysicalShape, const uint32_t* idxStrides) {
     this->srcLength = srcLength;
     this->coreTailDataNum = coreTailDataNum;
     this->blockLength = coreDataNum + (GetBlockNum() == GetBlockIdx() + 1 ? coreTailDataNum : 0);
     this->tileLength = ubPartDataNum;
     this->K = K;
     this->ALIGN_NUM = ALIGN_NUM;
     this->nSegments = nSegments;
     this->E_1 = E_1;
     this->indexPhysicalNum = indexPhysicalNum;
     this->aivNum = aivNum;
     this->logicalM = logicalM;
     this->logicalIndexNum = logicalIndexNum;
     this->strideE1 = strideE1;
     this->strideM = strideM;
     this->idxDimNum = idxDimNum;
      for (uint32_t d = 0; d < 8 && d < idxDimNum; ++d) {
         this->logicalShape[d] = logicalShape[d];
         this->idxPhysicalShape[d] = idxPhysicalShape[d];
         this->idxStrides[d] = idxStrides[d];
     }
     for (uint32_t d = idxDimNum; d < 8; ++d) {
         this->logicalShape[d] = 0;
         this->idxPhysicalShape[d] = 0;
         this->idxStrides[d] = 0;
     }
     auto startPointer = coreDataNum * GetBlockIdx();
     auto bufferlength = this->blockLength;
     
     auto outputlength = this->E_1 * this->nSegments * this->K;
     srcGm.SetGlobalBuffer((__gm__ T*)src, srcLength);
     indexGm.SetGlobalBuffer((__gm__ int32_t*)index, indexPhysicalNum);
     outGm.SetGlobalBuffer((__gm__ T*)out, outputlength);
     int32_t rows = static_cast<int32_t>(this->E_1 * this->nSegments);
     // Workspace holds per-row non-empty flags
     nonEmptyGm.SetGlobalBuffer((__gm__ uint8_t*)workspace, rows);
     this->tileNum = this->K / this->tileLength + (this->K %this->tileLength > 0);
     pipe.InitBuffer(inSrcQ, BUFFER_NUM, this->tileLength * sizeof(T));
     pipe.InitBuffer(ioMinQ, BUFFER_NUM, this->tileLength * sizeof(T));
     pipe.InitBuffer(idxQ, BUFFER_NUM, this->tileLength * sizeof(int32_t));
     if(GetBlockIdx() == 0){
         InitOutGmBuffers();
         InitNonEmptyGmBuffers();
     }
     SyncAll();
     }
 
 
     // Map linear index to multi-dimensional offset (supports broadcast)
     __aicore__ inline uint32_t ComputeIndexOffset(uint32_t lin) {
         if (this->idxDimNum <= 2) {
             uint32_t e = lin / this->logicalM;
             uint32_t m = lin - e * this->logicalM;
             return e * this->strideE1 + m * this->strideM;
         } else {
             uint32_t offset = 0;
             uint32_t remaining = lin;
             
             for (int32_t d = static_cast<int32_t>(this->idxDimNum) - 1; d >= 0; --d) {
                 uint32_t coord = remaining % this->logicalShape[d];
                 remaining = remaining / this->logicalShape[d];
                 offset += coord * this->idxStrides[d];
             }
             return offset;
         }
     }
     
     __aicore__ inline void Process(){
         uint32_t coreBeg = GetBlockIdx() * (this->blockLength);
         uint32_t coreEnd = coreBeg + this->blockLength;
         if (coreEnd > this->logicalIndexNum) coreEnd = this->logicalIndexNum;
         for(uint32_t lin = coreBeg; lin < coreEnd; ++lin){
             uint32_t idxOffset = ComputeIndexOffset(lin);
             int32_t p = indexGm.GetValue(idxOffset);
             
             uint32_t e = lin / this->logicalM;
             if (p < 0 || p >= static_cast<int32_t>(this->nSegments)) {
                 continue;
             }
 
             int32_t loopCount = this->tileNum;
             for(int32_t i=0; i<loopCount-1; i++){
                 CopyIn(static_cast<int32_t>(p), static_cast<int32_t>(lin), i, this->tileLength, 0);
                 Compute(i,this->tileLength);
                 CopyOut(static_cast<int32_t>(e), static_cast<int32_t>(p), i, this->tileLength);
             }    
             auto padding = (this->K % this->ALIGN_NUM ? this->ALIGN_NUM - this->K % this->ALIGN_NUM : 0);
             auto length = this->K + padding;
             length = length - this->tileLength * (loopCount - 1);
             CopyIn(static_cast<int32_t>(p), static_cast<int32_t>(lin), loopCount-1, length, padding);
             Compute(loopCount-1,length);
             CopyOut(static_cast<int32_t>(e), static_cast<int32_t>(p), loopCount-1, length);
             LocalTensor<uint8_t> oneFlag = idxQ.AllocTensor<uint8_t>();
             oneFlag.SetValue(0, static_cast<uint8_t>(1));
             int32_t seg = static_cast<int32_t>(p);
             int32_t row = static_cast<int32_t>(e) * static_cast<int32_t>(this->nSegments) + seg;
             GlobalTensor<uint8_t> flagDst;
             flagDst.SetGlobalBuffer((__gm__ uint8_t*)nonEmptyGm.GetPhyAddr() + row, 1);
             DataCopyExtParams copyParams{1, static_cast<uint32_t>(sizeof(uint8_t)), 0, 0, 0};
             DataCopyPad(flagDst, oneFlag, copyParams);
             idxQ.FreeTensor(oneFlag);
         }
 
             SyncAll();
             if(GetBlockIdx() == 0){
                 ProcessEmptySegments();
             }
             SyncAll();
         
         
     }
 private:
     __aicore__ inline void  InitOutGmBuffers(){
         int32_t n_ = this->nSegments * this->E_1;
         LocalTensor<T> tileBuffer = ioMinQ.AllocTensor<T>();
         for(int32_t i = 0; i < this->tileLength; i++){
             tileBuffer.SetValue(i, SegmentMinTraits<T>::InitValue());
         }
         for(int32_t seg = 0; seg < n_; seg++ ){
             for(int32_t tile = 0; tile < this->K; tile += this->tileLength){
                 GlobalTensor<T> segGm;
                 segGm.SetGlobalBuffer((__gm__ T*)outGm.GetPhyAddr() + seg * this->K + tile, this->tileLength);
                 if(tile + this->tileLength <= this->K){
                     DataCopy(segGm,tileBuffer,this->tileLength);
                 }else{
                     DataCopyExtParams copyParams{1, static_cast<uint32_t>((this->K % this->tileLength) * sizeof(T)), 0, 0, 0};
                     DataCopyPad(segGm,tileBuffer,copyParams);
                 }
             }
         }
         ioMinQ.FreeTensor(tileBuffer);
         
     }
     __aicore__ inline void InitNonEmptyGmBuffers(){
         // Initialize non-empty flags to 0 (Duplicate does not support uint8_t)
         int32_t rows = static_cast<int32_t>(this->E_1 * this->nSegments);
         LocalTensor<uint8_t> zeroFlag = idxQ.AllocTensor<uint8_t>();
         int32_t flagTile = static_cast<int32_t>(this->tileLength);
         for(int32_t i = 0; i < flagTile; i++){
             zeroFlag.SetValue(i, static_cast<uint8_t>(0));
         }
         for(int32_t t = 0; t < rows; t += flagTile){
             GlobalTensor<uint8_t> flagGm;
             flagGm.SetGlobalBuffer((__gm__ uint8_t*)nonEmptyGm.GetPhyAddr() + t, flagTile);
             if(t + flagTile <= rows){
                 DataCopy(flagGm, zeroFlag, flagTile);
             }else{
                 DataCopyExtParams copyParams{1, static_cast<uint32_t>((rows % flagTile) * sizeof(uint8_t)), 0, 0, 0};
                 DataCopyPad(flagGm, zeroFlag, copyParams);
             }
         }
         idxQ.FreeTensor(zeroFlag);
     }
 
     __aicore__ inline void ProcessEmptySegments(){
         LocalTensor<T> zero = ioMinQ.AllocTensor<T>();
         Duplicate(zero, SegmentMinTraits<T>::ZeroValue(), this->tileLength);
         
         int32_t rows = static_cast<int32_t>(this->nSegments * this->E_1);
         for(int32_t row = 0; row < rows; row++ ){
             if(nonEmptyGm.GetValue(row) == static_cast<uint8_t>(0)){
                 for(int32_t tile = 0; tile < static_cast<int32_t>(this->K); tile += this->tileLength){
                     GlobalTensor<T> segGm;
                     segGm.SetGlobalBuffer((__gm__ T*)outGm.GetPhyAddr() + row * this->K + tile, this->tileLength);
                     if(tile + this->tileLength <= static_cast<int32_t>(this->K)){
                         DataCopy(segGm, zero, this->tileLength);
                     }else{
                         DataCopyExtParams copyParams{1, static_cast<uint32_t>((this->K % this->tileLength) * sizeof(T)), 0, 0, 0};
                         DataCopyPad(segGm, zero, copyParams);
                     }
                 }
             }
         }
 
         ioMinQ.FreeTensor(zero);
     }
     
 
     __aicore__ inline void CopyIn(int32_t p, int32_t globalJ, int32_t progress, uint32_t length, uint32_t padding){
         LocalTensor<T> dLocal = inSrcQ.AllocTensor<T>();
         DataCopy(dLocal, srcGm[globalJ * this->K + progress * this->tileLength], length);
         for(int i=length-padding;i<length;i++){
             dLocal.SetValue(i, SegmentMinTraits<T>::InitValue());
         }
         inSrcQ.EnQue(dLocal);
     }
     __aicore__ inline void Compute(int32_t progress, int32_t length){
         LocalTensor<T> dLocal = inSrcQ.DeQue<T>();
         LocalTensor<T> zLocal = ioMinQ.AllocTensor<T>();
 
         DataCopy(zLocal, dLocal, length);
         ioMinQ.EnQue(zLocal);
         inSrcQ.FreeTensor(dLocal);
 
     }
     __aicore__ inline void CopyOut(int32_t e, int32_t seg, int32_t progress, uint32_t length){
         LocalTensor<T> zLocal = ioMinQ.DeQue<T>();
         SetAtomicMin<T>();
         // Output layout is [E, nSegments, K]; offset = e * nSegments * K + seg * K + progress * tileLength
         uint32_t outOffset = e * this->nSegments * this->K + seg * this->K + progress * this->tileLength;
         DataCopy(outGm[outOffset], zLocal, length);
         SetAtomicNone();
 
         ioMinQ.FreeTensor(zLocal);
     }
 private:
     TPipe pipe;
     TQue<TPosition::VECIN, BUFFER_NUM> inSrcQ;
     TQue<TPosition::VECOUT, BUFFER_NUM> ioMinQ;
     TQue<TPosition::VECIN,BUFFER_NUM> idxQ;
     GlobalTensor<T> srcGm;
     GlobalTensor<int32_t> indexGm;
     GlobalTensor<T> outGm;
     GlobalTensor<uint8_t> nonEmptyGm;
     uint32_t srcLength;
     uint32_t coreTailDataNum;
     uint32_t blockLength;
     uint32_t tileLength;
     uint32_t K;
     uint32_t ALIGN_NUM;
     uint32_t nSegments;
     uint32_t E_1;
     uint32_t tileNum;
     uint32_t indexPhysicalNum;
     uint32_t aivNum;
     uint32_t logicalM;
     uint32_t logicalIndexNum;
     uint32_t strideE1;
     uint32_t strideM;
     uint32_t idxDimNum;
     uint32_t logicalShape[8];
     uint32_t idxPhysicalShape[8];
     uint32_t idxStrides[8];
 };
 
 template <typename T>
 class SegmentMinCooArgKernel {
 public:
     __aicore__ inline SegmentMinCooArgKernel(){}
     __aicore__ inline void Init(GM_ADDR src, GM_ADDR index, GM_ADDR out, GM_ADDR arg_out,
                                 uint32_t srcLength, uint32_t ubPartDataNum,
                                 uint32_t E_1, uint32_t K, uint32_t coreDataNum, uint32_t coreTailDataNum,
                                 uint32_t nSegments, uint32_t ALIGN_NUM, uint32_t indexPhysicalNum, uint32_t aivNum,
                                 uint32_t logicalM, uint32_t logicalIndexNum, uint32_t strideE1, uint32_t strideM,
                                 uint32_t idxDimNum, const uint32_t* logicalShape, const uint32_t* idxPhysicalShape, const uint32_t* idxStrides) {
         this->coreTailDataNum = coreTailDataNum;
         this->blockLength = coreDataNum + (GetBlockNum() == GetBlockIdx() + 1 ? coreTailDataNum : 0);
         this->tileLength = ubPartDataNum;
         this->K = K;
         this->ALIGN_NUM = ALIGN_NUM;
         this->nSegments = nSegments;
         this->E_1 = E_1;
         this->indexPhysicalNum = indexPhysicalNum;
         this->aivNum = aivNum;
         this->logicalM = logicalM;
         this->logicalIndexNum = logicalIndexNum;
         this->strideE1 = strideE1;
         this->strideM = strideM;
         this->idxDimNum = idxDimNum;
         for (uint32_t d = 0; d < 8 && d < idxDimNum; ++d) {
             this->logicalShape[d] = logicalShape[d];
             this->idxPhysicalShape[d] = idxPhysicalShape[d];
             this->idxStrides[d] = idxStrides[d];
         }
         for (uint32_t d = idxDimNum; d < 8; ++d) {
             this->logicalShape[d] = 0;
             this->idxPhysicalShape[d] = 0;
             this->idxStrides[d] = 0;
         }
         auto startPointer = coreDataNum * GetBlockIdx();
         auto bufferlength = this->blockLength;
         
         int32_t M = static_cast<int32_t>(this->logicalM);
         auto outputlength = this->E_1 * M * this->nSegments;
         srcGm.SetGlobalBuffer((__gm__ T*)src, srcLength);
         indexGm.SetGlobalBuffer((__gm__ int32_t*)index, indexPhysicalNum);
         outGm.SetGlobalBuffer((__gm__ T*)out, outputlength);
         argOutGm.SetGlobalBuffer((__gm__ int32_t*)arg_out, outputlength);
         
         this->tileNum = this->K / this->tileLength + (this->K % this->tileLength > 0);
         pipe.InitBuffer(inSrcQ, BUFFER_NUM, this->tileLength * sizeof(T));
         pipe.InitBuffer(inOutQ, BUFFER_NUM, this->tileLength * sizeof(T));
         pipe.InitBuffer(idxQ, BUFFER_NUM, this->tileLength * sizeof(int32_t));
         
         if(GetBlockIdx() == 0){
             InitArgOutBuffers();
         }
         SyncAll();
     }
     
     // Map linear index to multi-dimensional offset (supports broadcast)
     __aicore__ inline uint32_t ComputeIndexOffset(uint32_t lin) {
         if (this->idxDimNum <= 2) {
             uint32_t e = lin / this->logicalM;
             uint32_t m = lin - e * this->logicalM;
             return e * this->strideE1 + m * this->strideM;
         } else {
             uint32_t offset = 0;
             uint32_t remaining = lin;
             
             for (int32_t d = static_cast<int32_t>(this->idxDimNum) - 1; d >= 0; --d) {
                 uint32_t coord = remaining % this->logicalShape[d];
                 remaining = remaining / this->logicalShape[d];
                 offset += coord * this->idxStrides[d];
             }
             return offset;
         }
     }
     
     __aicore__ inline void Process(){
        // 使用 LocalTensor 缓存数据，避免直接访问 GlobalTensor 的缓存一致性问题
        LocalTensor<T> srcLocal = inSrcQ.AllocTensor<T>();
        LocalTensor<T> outLocal = inOutQ.AllocTensor<T>();
        LocalTensor<int32_t> argLocal = idxQ.AllocTensor<int32_t>();
        
        uint32_t coreBeg = GetBlockIdx() * (this->blockLength);
        uint32_t coreEnd = coreBeg + this->blockLength;
        if (coreEnd > this->logicalIndexNum) coreEnd = this->logicalIndexNum;
        for(uint32_t lin = coreBeg; lin < coreEnd; ++lin){
            // 使用多维索引计算（自动兼容2D情况）
            uint32_t idxOffset = ComputeIndexOffset(lin);
            int32_t seg = indexGm.GetValue(idxOffset);
            
            // 为了兼容输出布局，仍需要计算 e_1（用于输出索引）
            int32_t e_1 = static_cast<int32_t>(lin / this->logicalM);
            int32_t E_2 = static_cast<int32_t>(this->logicalM);
            // 对于多维情况，e_2 需要从多维坐标中提取
            // 这里简化处理：e_2 = lin % logicalM（对于2D情况正确，对于多维需要更复杂的分解）
            int32_t e_2 = static_cast<int32_t>(lin - static_cast<uint32_t>(e_1) * this->logicalM);
            for(int32_t k = 0; k < this->K; k++){
                uint32_t outOffset = e_1 * this->nSegments * this->K + seg * this->K + k;
                uint32_t srcOffset = e_1 * E_2 * this->K + e_2 * this->K + k;
                uint32_t argOutOffset = e_1 * this->nSegments * this->K + seg * this->K + k;
                // 读取 src 值（单个元素）
                srcLocal.SetValue(0, srcGm.GetValue(srcOffset));
                
                // 读取 out 和 arg_out 的值（单个元素）
                outLocal.SetValue(0, outGm.GetValue(outOffset));
                argLocal.SetValue(0, argOutGm.GetValue(argOutOffset));
                
                // 转换为 float 进行比较（使用绝对容差，与 gen_data.py 保持一致）
                T srcVal = srcLocal.GetValue(0);
                T outVal = outLocal.GetValue(0);
                int32_t currentArg = argLocal.GetValue(0);
                
                bool isEqual = SegmentMinTraits<T>::IsClose(srcVal, outVal);
                
                // 如果 src 等于 out，且 arg_out 还是 -1，则写入 k
                if(isEqual){
                    argLocal.SetValue(0, e_2);
                    // 写回 arg_out（使用 DataCopyPad 确保正确写入）
                    DataCopyExtParams copyParams{1, static_cast<uint32_t>(sizeof(int32_t)), 0, 0, 0};
                    DataCopyPad(argOutGm[argOutOffset], argLocal, copyParams);
            }
            }

        
            
            
        }
        
        inSrcQ.FreeTensor(srcLocal);
        inOutQ.FreeTensor(outLocal);
        idxQ.FreeTensor(argLocal);
    }
     
 private:
     __aicore__ inline void InitArgOutBuffers(){
         int32_t M = static_cast<int32_t>(this->logicalM);
         int32_t totalSize = this->E_1 * this->nSegments * this->K;
         LocalTensor<int32_t> negOneBuffer = idxQ.AllocTensor<int32_t>();
         for(int32_t i = 0; i < this->tileLength; i++){
             negOneBuffer.SetValue(i, static_cast<int32_t>(this->nSegments));
         }
         
         for(int32_t offset = 0; offset < totalSize; offset += this->tileLength){
             GlobalTensor<int32_t> argGm;
             int32_t remaining = totalSize - offset;
             int32_t copyLength = (remaining >= static_cast<int32_t>(this->tileLength)) ? this->tileLength : remaining;
             argGm.SetGlobalBuffer((__gm__ int32_t*)argOutGm.GetPhyAddr() + offset, copyLength);
             if(copyLength == static_cast<int32_t>(this->tileLength)){
                 DataCopy(argGm, negOneBuffer, this->tileLength);
             }else{
                 DataCopyExtParams copyParams{1, static_cast<uint32_t>(copyLength * sizeof(int32_t)), 0, 0, 0};
                 DataCopyPad(argGm, negOneBuffer, copyParams);
             }
         }
         idxQ.FreeTensor(negOneBuffer);
     }
     
     
 private:
     TPipe pipe;
     TQue<TPosition::VECIN, BUFFER_NUM> inSrcQ;
     TQue<TPosition::VECIN, BUFFER_NUM> inOutQ;
     TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> idxQ;
     GlobalTensor<T> srcGm;
     GlobalTensor<int32_t> indexGm;
     GlobalTensor<T> outGm;
     GlobalTensor<int32_t> argOutGm;
     uint32_t coreTailDataNum;
     uint32_t blockLength;
     uint32_t tileLength;
     uint32_t K;
     uint32_t ALIGN_NUM;
     uint32_t nSegments;
     uint32_t E_1;
     uint32_t tileNum;
     uint32_t indexPhysicalNum;
     uint32_t aivNum;
     uint32_t logicalM;
     uint32_t logicalIndexNum;
     uint32_t strideE1;
     uint32_t strideM;
     uint32_t idxDimNum;
     uint32_t logicalShape[8];
     uint32_t idxPhysicalShape[8];
     uint32_t idxStrides[8];
 };
 
 extern "C" __global__ __aicore__ void segment_min_coo(GM_ADDR src, GM_ADDR index, GM_ADDR out, GM_ADDR arg_out, GM_ADDR workspace, GM_ADDR tiling) {
     GET_TILING_DATA(tiling_data, tiling);
     TPipe tpipe;
     
     uint32_t idxDimNum = tiling_data.idxDimNum;
     uint32_t logicalShape[8];
     uint32_t idxPhysicalShape[8];
     uint32_t idxStrides[8];
     for (uint32_t d = 0; d < 8; ++d) {
         logicalShape[d] = tiling_data.logicalShape[d];
         idxPhysicalShape[d] = tiling_data.idxPhysicalShape[d];
         idxStrides[d] = tiling_data.idxStrides[d];
     }
     
     // ========== Pass 1: compute minimum ==========
     SegmentMinCooKernel<SrcDataType> op1;
     op1.Init(src, index, out, workspace, tiling_data.srcLength, tiling_data.ubPartDataNum,
         tiling_data.E_1, tiling_data.K, tiling_data.coreDataNum, tiling_data.coreTailDataNum, 
         tiling_data.nSegments, tiling_data.ALIGN_NUM, tiling_data.indexPhysicalNum, tiling_data.aivNum,
         tiling_data.logicalM, tiling_data.logicalIndexNum, tiling_data.strideE1, tiling_data.strideM,
         idxDimNum, logicalShape, idxPhysicalShape, idxStrides);
     op1.Process();
     
     tpipe.Reset();
     SyncAll();
     
     // ========== Pass 2: compute arg_out ==========
     SegmentMinCooArgKernel<SrcDataType> op2;
     op2.Init(src, index, out, arg_out, tiling_data.srcLength, tiling_data.ubPartDataNum,
         tiling_data.E_1, tiling_data.K, tiling_data.coreDataNum, tiling_data.coreTailDataNum,
         tiling_data.nSegments, tiling_data.ALIGN_NUM, tiling_data.indexPhysicalNum, tiling_data.aivNum,
         tiling_data.logicalM, tiling_data.logicalIndexNum, tiling_data.strideE1, tiling_data.strideM,
         idxDimNum, logicalShape, idxPhysicalShape, idxStrides);
     op2.Process();
 }